package com.example.base;

import com.example.entity.Order;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.streaming.api.datastream.AllWindowedStream;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.datastream.WindowedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessAllWindowFunction;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;
import org.apache.flink.util.OutputTag;

import java.time.Duration;

/**
 * Author wangJinLong
 * Date 2025/8/13 17:33
 **/
@Slf4j
public class FlinkStreamWindowFunction {
    public static void main(String[] args) throws Exception {
        outPutTag();
    }

    private static void outPutTag() throws Exception {
        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();
        environment.setParallelism(1);

        DataStream<Order> mapSource = FlinkStreamSource.getMapSource(environment);

        OutputTag<Order> orderOutputTag = new OutputTag<>("order-late-data"){};

        WindowedStream<Order, String, TimeWindow> windowedStream =
                mapSource.keyBy(Order::getOrderId)
                        .window(TumblingEventTimeWindows.of(Duration.ofSeconds(10)))
                        .allowedLateness(Duration.ofSeconds(2))
                        .sideOutputLateData(orderOutputTag);

        SingleOutputStreamOperator<Order> process = windowedStream.process(new ProcessWindowFunction<Order, Order, String, TimeWindow>() {
            @Override
            public void process(String s, Context context, Iterable<Order> elements, Collector<Order> out) throws Exception {
                elements.forEach(out::collect);
            }
        });

        process.getSideOutput(orderOutputTag).print("late-data");

        environment.execute();
    }

    private static void aggrateWindow() throws Exception {
        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();
        environment.setParallelism(3);

        DataStream<Order> mapSource = FlinkStreamSource.getMapSource(environment);

        WindowedStream<Order, String, TimeWindow> windowedStream = mapSource.keyBy(Order::getOrderId)
                .window(TumblingEventTimeWindows.of(Duration.ofSeconds(10)))
                .allowedLateness(Duration.ofSeconds(2))
                ;

        SingleOutputStreamOperator<Order> aggregateStream = windowedStream.aggregate(new AggregateFunction<Order, Order, Order>() {
            @Override
            public Order createAccumulator() {
                log.info("createAccumulator...");
                Order order = new Order();
                order.setAmount(1);
                return order;
            }

            @Override
            public Order add(Order value, Order accumulator) {
                log.info("add...");
                Order order = new Order();
                order.setOrderId(value.getOrderId());
                order.setAmount(value.getAmount() * accumulator.getAmount());
                order.setTotalPrice(value.getTotalPrice());
                order.setOrderTime(value.getOrderTime());
                return order;
            }

            @Override
            public Order getResult(Order accumulator) {
                log.info("getResult...");
                return accumulator;
            }

            @Override
            public Order merge(Order a, Order b) {
                Order order = new Order();
                order.setOrderId(a.getOrderId());
                order.setAmount(a.getAmount() + b.getAmount());
                log.info("merge...");
                return order;
            }
        });

        aggregateStream.print();

        environment.execute();
    }

    private static void processWindow() throws Exception {
        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();
        environment.setParallelism(1);

        DataStream<Order> mapSource = FlinkStreamSource.getMapSource(environment);

        AllWindowedStream<Order, TimeWindow> windowAll = mapSource.windowAll(TumblingEventTimeWindows.of(Duration.ofSeconds(10)));

        SingleOutputStreamOperator<Integer> processStream = windowAll.process(new ProcessAllWindowFunction<Order, Integer, TimeWindow>() {
            @Override
            public void process(Context context, Iterable<Order> elements, Collector<Integer> out) throws Exception {
                int num = 0;
                for (Order element : elements) {
                    num += element.getAmount();
                }
                out.collect(num);
            }
        });

        processStream.print();

        environment.execute();
    }

    private static void reduceWindow() throws Exception {
        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();
        environment.setParallelism(1);

        DataStream<Order> mapSource = FlinkStreamSource.getMapSource(environment);

//        mapSource.print();

        AllWindowedStream<Order, TimeWindow> windowAll = mapSource.windowAll(TumblingEventTimeWindows.of(Duration.ofSeconds(10)));

        SingleOutputStreamOperator<Order> reduceStream = windowAll.reduce(new ReduceFunction<Order>() {

            @Override
            public Order reduce(Order value1, Order value2) throws Exception {
                value1.setAmount(value1.getAmount() * value2.getAmount());
                return value1;
            }
        });

        reduceStream.print();

        environment.execute();
    }
}
